#pragma once

#if (defined(__MINGW32__) || defined(_MSC_VER)) && !defined(_WIN32_WINNT)
#define _WIN32_WINNT 0x0601
#endif

#include <string>
#include <set>
#include <functional>
#include <memory>				//shared_ptr use
#include <queue>
#ifndef ASIO_STANDALONE
#define ASIO_STANDALONE
#define ASIO_NO_DEPRECATED
#endif
#include "asio/io_context.hpp" 	//io_context use
#include "asio/ip/udp.hpp"		//acceptor and socket use
#include "asiobuf.h"

namespace common
{

template<typename HEAD, typename USERDATA>
class UdpServer;

template<typename HEAD, typename USERDATA = bool>
class UdpSession : public std::enable_shared_from_this<UdpSession<HEAD, USERDATA>>
{
public:
	UdpSession(UdpServer<HEAD, USERDATA> &srv, asio::ip::udp::endpoint _ep) : server(srv), ep(_ep) {}
	~UdpSession() { server.onClose(ip, userdata); }

	const std::string & address()
	{
		if (ip.empty()) {
			ip = std::move(ep.address().to_string());
		}
		return ip;
	}

	//添加发送消息
	void add_sndmsg(std::shared_ptr<common::Buffer> buffer)
	{
		if (!buffer || !(*buffer))
			return;
		using SessionSPtr = typename UdpServer<HEAD, USERDATA>::SessionSPtr;
		server.add_sndmsg(std::pair<SessionSPtr, std::shared_ptr<common::Buffer>>(this->shared_from_this(), std::move(buffer)));
	}
	//获取session是否处于保活状态
	bool keep() const {return keepFlag;}
	//设置session是否处于保活状态,如果设置为false则会在onSend之后销毁session
	void keep(bool yes) {keepFlag = yes;}

private:
	UdpServer<HEAD, USERDATA> &server;
	std::string ip;
	bool keepFlag{false};	//保活标志,如果为false,则会在onSend之后销毁session
public:
	asio::ip::udp::endpoint ep;
	USERDATA userdata;
};

template<typename HEAD, typename USERDATA = bool>
class UdpServer
{
public:
	using UserData = USERDATA;
	using SessionSPtr = std::shared_ptr<UdpSession<HEAD, USERDATA>>;
	using SessionWPtr = std::weak_ptr<UdpSession<HEAD, USERDATA>>;
	struct SessionSPtrCMP {  
		bool operator() (const SessionSPtr& lhs, const SessionSPtr& rhs) const  
		{  
			return lhs->ep < rhs->ep;  
		} 
	};
	using SessionSet = std::set<SessionSPtr, SessionSPtrCMP>;
	using iterator = typename SessionSet::iterator;
	using const_iterator = typename SessionSet::const_iterator;

	UdpServer() : socket(m_ioctx) {}

	virtual ~UdpServer() 
	{
		asio::error_code ec;
		socket.shutdown(asio::ip::tcp::socket::shutdown_both, ec);
		socket.close(ec);
		m_ioctx.stop();
	}
	
	//启动服务器,会阻塞
	void start(uint16_t port)
	{
		using asio::ip::udp;
		socket = udp::socket(m_ioctx, udp::endpoint(udp::v4(), port));
		do_receive();
		m_ioctx.run();
	}

	//添加发送信息到发送队列
	void add_sndmsg(std::pair<SessionSPtr, std::shared_ptr<common::Buffer>> &&data)
	{
		sndque.emplace(std::move(data));
		if (sndque.size() == 1) {
			do_send();
		}
	}
	//销毁客户端session
	void destroy(SessionSPtr pSession) {m_sessions.erase(pSession);}

	//对服务器所连接的所有客户端发布消息
	void publish(std::shared_ptr<common::Buffer> buffer)
	{
		if (!buffer || !(*buffer))
			return;
		auto callback = [buffer, this]{
			for (auto &it : m_sessions) {
				it->add_sndmsg(buffer);
			}
		};
		asio::post(m_ioctx, callback);
	}
	//通知其他客户端，这个主要用来对其他客户端发布配置更新消息
	void notify(std::shared_ptr<common::Buffer> buffer, SessionSPtr pSession)
	{
		if (!buffer || !(*buffer))
			return;
		auto callback = [buffer, pSession, this]{
			for (auto &it : m_sessions) {
				if (it != pSession)
					it->add_sndmsg(buffer);
			}
		};
		asio::post(m_ioctx, callback);
	}
	//发布一个任务交由UdpServer的ioctx执行
	void post(std::function<void()> task)
	{
		asio::post(m_ioctx, std::move(task));
	}

	//客户端连接之后调用
	virtual void onConnect(SessionSPtr pSession) {}
	//读取到消息头部之后调用,返回后面要接收的消息体size(即去除HEAD),返回-1表示接收的head有误应关闭session,返回0表示本次数据只有head没有body
	virtual int onHead(HEAD &head) {return 0;}
	//读取到客户端数据之后调用
	virtual bool onMessage(std::shared_ptr<common::Buffer> buffer, SessionSPtr pSession) = 0;
	//发送成功之后调用
	virtual void onSend(SessionSPtr pSession) {}
	//客户端关闭之后调用
	virtual void onClose(const std::string &ip, UserData &data) {}

	iterator begin() {return m_sessions.begin();}
	const_iterator begin() const {return m_sessions.begin();}
	iterator end() {return m_sessions.end();}
	const_iterator end() const {return m_sessions.end();}

private:
	asio::io_context 		m_ioctx;
	asio::ip::udp::socket 	socket;
	asio::ip::udp::endpoint ep;
	asio::RecvBuffer 		rcvbuf{65536};
	SessionSet 				m_sessions;
	std::queue<std::pair<SessionSPtr, std::shared_ptr<common::Buffer>>> sndque;

	void do_receive()
	{
		socket.async_receive_from(rcvbuf, ep, [this](std::error_code ec, std::size_t bytes_recvd){   
			if (!ec && bytes_recvd >= sizeof(HEAD)) {
				HEAD head;
				memcpy(&head, rcvbuf.data(), sizeof(HEAD));
				if (bytes_recvd >= sizeof(HEAD) + onHead(head)) {
					SessionSPtr sptr = nullptr;
					for (auto &it : m_sessions) {
						if (it->ep == ep) {
							sptr = it;
							break;
						}
					}
					if (sptr == nullptr) {
						sptr = std::make_shared<UdpSession<HEAD, USERDATA>>(*this, std::move(ep));
						m_sessions.insert(sptr);
						onConnect(sptr);
					}
					auto buf = std::make_shared<common::Buffer>(bytes_recvd);
					buf->write((const char *)rcvbuf.data(), bytes_recvd);
					if (!onMessage(std::move(buf), sptr)) {
						destroy(sptr);
					}
				}
			}  
			do_receive();
		}); 
	}

	void do_send()
	{
		if (sndque.empty())
			return;
		auto &data = sndque.front();
		auto &buf = data.second;
		socket.async_send_to(asio::SendBuffer(buf), data.first->ep, 
			[this](std::error_code ec, std::size_t bytes_sent){
				auto &pSession = sndque.front().first;
				if (!ec) {
					onSend(pSession);
					if (!pSession->keep())
						destroy(pSession);
				} else {
					destroy(pSession);
				}
				sndque.pop();  
	      		do_send();
	    }); 
	}
};

}

